#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import serial
import struct
import time
import sys
import os
import argparse
from intelhex import IntelHex

class UartUpgrade:
    # 协议常量定义 - 与固件端一致
    SOF = 0xAA                # 帧头
    # 命令定义
    CMD_START = 0x01          # 升级请求
    CMD_DATA = 0x02           # 数据传输
    CMD_FINISH = 0x03         # 升级完成
    CMD_VERIFY = 0x04         # 验证固件
    CMD_RESET = 0x05          # 系统复位
    CMD_ACK = 0x80            # 确认响应
    CMD_NACK = 0x81           # 拒绝响应
    
    # 固件参数
    BOOTLOADER_ADDR = 0x08000000  # Bootloader起始地址
    BOOTLOADER_SIZE = 0x8000      # Bootloader大小(32KB)
    APP_ADDR = 0x08008000         # APP起始地址 (Bootloader地址 + Bootloader大小)
    
    # 通信参数
    PACKET_SIZE = 256             # 每个数据包大小
    TIMEOUT = 3.0                 # 超时时间(秒)
    ACK_TIMEOUT = 1.0             # ACK超时时间(秒)
    RETRY_COUNT = 3               # 重试次数
    
    def __init__(self, port, baudrate=115200, timeout=1.0):
        """初始化串口升级实例"""
        self.serial = serial.Serial(port, baudrate, timeout=timeout)
        # 清空接收缓冲区
        self.serial.reset_input_buffer()
        # 等待串口稳定
        time.sleep(0.2)
        
    def __del__(self):
        """析构函数，关闭串口"""
        if hasattr(self, 'serial') and self.serial.is_open:
            self.serial.close()
            
    @staticmethod
    def calculate_crc16(data):
        """计算MODBUS CRC16 - 与固件使用相同算法"""
        crc = 0xFFFF
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ 0xA001
                else:
                    crc >>= 1
        return crc
    
    def send_frame(self, cmd, data=b''):
        """发送一帧数据，并自动重试"""
        # 构建数据帧：SOF(1) + CMD(1) + LEN(2) + DATA(n) + CRC(2)
        frame = struct.pack('<BBH', self.SOF, cmd, len(data)) + data
        crc = self.calculate_crc16(frame)
        frame += struct.pack('<H', crc)
        
        # 打印调试信息
        print(f"发送: CMD=0x{cmd:02X}, LEN={len(data)}")
        
        # 发送数据
        self.serial.write(frame)
        self.serial.flush()  # 确保数据发送完成
        
    def read_byte_with_timeout(self, timeout):
        """带超时的读取一个字节"""
        start_time = time.time()
        while (time.time() - start_time) < timeout:
            if self.serial.in_waiting > 0:
                return self.serial.read(1)
            time.sleep(0.01)  # 短暂睡眠，避免占用CPU
        return None
        
    def receive_frame(self, timeout=TIMEOUT):
        """接收一帧数据，带超时处理"""
        start_time = time.time()
        
        # 等待帧头
        while (time.time() - start_time) < timeout:
            byte = self.read_byte_with_timeout(0.1)
            if byte is None:
                continue
            if byte == bytes([self.SOF]):
                break
        else:
            raise TimeoutError("等待帧头超时")
                
        # 读取命令和长度
        header_data = self.serial.read(3)
        if len(header_data) < 3:
            raise TimeoutError("读取命令和长度超时")
            
        cmd, length = struct.unpack('<BH', header_data)
        
        # 读取数据
        data = self.serial.read(length)
        if len(data) < length:
            raise TimeoutError("读取数据超时")
            
        # 读取CRC
        crc_data = self.serial.read(2)
        if len(crc_data) < 2:
            raise TimeoutError("读取CRC超时")
            
        crc_received = struct.unpack('<H', crc_data)[0]
        
        # 验证CRC
        frame = struct.pack('<BBH', self.SOF, cmd, length) + data
        crc_calc = self.calculate_crc16(frame)
        
        if crc_calc != crc_received:
            raise ValueError(f"CRC校验错误: 计算={crc_calc:04X}, 接收={crc_received:04X}")
        
        # 打印调试信息    
        print(f"接收: CMD=0x{cmd:02X}, LEN={length}")
            
        return cmd, data
        
    def wait_ack(self, timeout=ACK_TIMEOUT):
        """等待ACK响应，带超时处理"""
        try:
            cmd, _ = self.receive_frame(timeout)
            if cmd == self.CMD_NACK:
                raise RuntimeError("设备拒绝请求(NACK)")
            elif cmd != self.CMD_ACK:
                raise RuntimeError(f"收到未知响应: 0x{cmd:02X}")
            return True
        except Exception as e:
            print(f"等待ACK失败: {str(e)}")
            return False
            
    def retry_command(self, func, retries=RETRY_COUNT):
        """自动重试命令执行"""
        for attempt in range(retries):
            try:
                return func()
            except Exception as e:
                if attempt < retries - 1:
                    print(f"命令执行失败，正在重试({attempt+1}/{retries}): {str(e)}")
                    time.sleep(0.5)  # 重试前等待一段时间
                else:
                    raise
    
    def enter_bootloader(self):
        """发送命令让设备进入bootloader模式"""
        print("正在请求设备进入bootloader模式...")
        
        # 发送升级请求命令
        self.send_frame(self.CMD_START)
        
        # 等待设备响应
        if self.wait_ack():
            print("设备已进入bootloader模式")
            # 设备进入bootloader后会复位，等待一段时间再继续
            time.sleep(1.0)
            return True
        else:
            print("进入bootloader模式失败")
            return False
        
    def program_firmware(self, hex_file, verify=True):
        """编程固件主流程"""
        print("正在读取HEX文件...")
        ih = IntelHex(hex_file)
        
        # 获取固件数据
        firmware_data = ih.tobinarray(start=self.APP_ADDR)
        total_size = len(firmware_data)
        
        print(f"固件大小: {total_size} 字节")
        
        # 尝试进入bootloader模式
        if not self.retry_command(self.enter_bootloader):
            print("无法进入bootloader模式，请检查设备连接")
            return False
        
        # 发送升级请求与固件大小
        print("发送升级开始请求...")
        self.send_frame(self.CMD_START, struct.pack('<I', total_size))
        if not self.wait_ack():
            print("启动升级失败，设备未响应")
            return False
        
        # 分包发送数据
        sent_size = 0
        while sent_size < total_size:
            chunk_size = min(self.PACKET_SIZE, total_size - sent_size)
            chunk = firmware_data[sent_size:sent_size + chunk_size]
            
            # 显示进度
            progress = int(sent_size / total_size * 50)
            print(f"\r进度: [{'#' * progress}{' ' * (50-progress)}] {sent_size}/{total_size} ({sent_size/total_size*100:.1f}%)", end='')
            
            # 发送数据包并带偏移量
            packet_data = struct.pack('<I', sent_size) + chunk
            self.send_frame(self.CMD_DATA, packet_data)
            
            # 等待ACK响应
            if not self.wait_ack():
                print(f"\n发送数据包失败，偏移={sent_size}")
                return False
            
            sent_size += chunk_size
            
        print(f"\r进度: [{'#' * 50}] {total_size}/{total_size} (100%)")
        
        # 发送完成命令
        print("发送升级完成命令...")
        self.send_frame(self.CMD_FINISH)
        if not self.wait_ack():
            print("升级完成确认失败")
            return False
        
        # 验证固件
        if verify:
            print("正在验证固件...")
            self.send_frame(self.CMD_VERIFY)
            if not self.wait_ack():
                print("固件验证失败")
                return False
            print("固件验证成功")
            
        # 发送复位命令
        print("发送系统复位命令...")
        self.send_frame(self.CMD_RESET)
        
        print("固件升级完成！")
        return True

def main():
    """主函数"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='串口固件升级工具')
    parser.add_argument('port', help='串口号，如COM3或/dev/ttyUSB0')
    parser.add_argument('hex_file', help='HEX格式固件文件路径')
    parser.add_argument('-b', '--baudrate', type=int, default=115200, help='波特率(默认115200)')
    parser.add_argument('-v', '--verify', action='store_true', help='升级后验证固件')
    parser.add_argument('-r', '--retry', type=int, default=3, help='失败重试次数')
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.hex_file):
        print(f"错误: 找不到文件 {args.hex_file}")
        return 1
        
    try:
        # 创建升级实例
        upgrader = UartUpgrade(args.port, args.baudrate)
        UartUpgrade.RETRY_COUNT = args.retry
        
        # 执行升级
        if upgrader.program_firmware(args.hex_file, args.verify):
            return 0
        else:
            return 1
            
    except KeyboardInterrupt:
        print("\n操作已取消")
        return 130
    except Exception as e:
        print(f"\n错误: {str(e)}")
        return 1

if __name__ == '__main__':
    sys.exit(main()) 